import json
from pathlib import Path

from langchain.document_loaders import JSONLoader


def normal_json_loader():
    file_path = "files/facebook_chat.json"
    data = json.loads(Path(file_path).read_text())
    print(data)


def json_loader():
    loader = JSONLoader(file_path="files/facebook_chat.json", text_content=False, jq_schema=".messages[].content")
    data = loader.load()
    print(data)


def metadata_func(record: dict, metadata: dict) -> dict:
    metadata["sender_name"] = record.get("sender_name")
    metadata["timestamp_ms"] = record.get("timestamp_ms")
    if "source" in metadata:
        source = metadata["source"].split("/")
        source = source[source.index("langchain-test"):]
        metadata["source"] = "/".join(source)

    return metadata


def json_loader_with_extracting_metadata():
    loader = JSONLoader(file_path="files/facebook_chat.json",
                        jq_schema=".messages[]",
                        text_content=False,
                        content_key="content",
                        metadata_func=metadata_func)
    data = loader.load()
    for item in data:
        print(item)


if __name__ == '__main__':
    json_loader_with_extracting_metadata()